#include "crypto.h"
#include <string.h>

// SHA256 常量
static const uint32_t K[64] = {
    0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5,
    0x3956c25b, 0x59f111f1, 0x923f82a4, 0xab1c5ed5,
    0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3,
    0x72be5d74, 0x80deb1fe, 0x9bdc06a7, 0xc19bf174,
    0xe49b69c1, 0xefbe4786, 0x0fc19dc6, 0x240ca1cc,
    0x2de92c6f, 0x4a7484aa, 0x5cb0a9dc, 0x76f988da,
    0x983e5152, 0xa831c66d, 0xb00327c8, 0xbf597fc7,
    0xc6e00bf3, 0xd5a79147, 0x06ca6351, 0x14292967,
    0x27b70a85, 0x2e1b2138, 0x4d2c6dfc, 0x53380d13,
    0x650a7354, 0x766a0abb, 0x81c2c92e, 0x92722c85,
    0xa2bfe8a1, 0xa81a664b, 0xc24b8b70, 0xc76c51a3,
    0xd192e819, 0xd6990624, 0xf40e3585, 0x106aa070,
    0x19a4c116, 0x1e376c08, 0x2748774c, 0x34b0bcb5,
    0x391c0cb3, 0x4ed8aa4a, 0x5b9cca4f, 0x682e6ff3,
    0x748f82ee, 0x78a5636f, 0x84c87814, 0x8cc70208,
    0x90befffa, 0xa4506ceb, 0xbef9a3f7, 0xc67178f2
};

// 初始哈希值
static const uint32_t H0[8] = {
    0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a,
    0x510e527f, 0x9b05688c, 0x1f83d9ab, 0x5be0cd19
};

// 辅助函数
static uint32_t rotr(uint32_t x, int n) {
    return (x >> n) | (x << (32 - n));
}

static uint32_t ch(uint32_t x, uint32_t y, uint32_t z) {
    return (x & y) ^ (~x & z);
}

static uint32_t maj(uint32_t x, uint32_t y, uint32_t z) {
    return (x & y) ^ (x & z) ^ (y & z);
}

static uint32_t sigma0(uint32_t x) {
    return rotr(x, 2) ^ rotr(x, 13) ^ rotr(x, 22);
}

static uint32_t sigma1(uint32_t x) {
    return rotr(x, 6) ^ rotr(x, 11) ^ rotr(x, 25);
}

static uint32_t gamma0(uint32_t x) {
    return rotr(x, 7) ^ rotr(x, 18) ^ (x >> 3);
}

static uint32_t gamma1(uint32_t x) {
    return rotr(x, 17) ^ rotr(x, 19) ^ (x >> 10);
}

// 字节序转换
static uint32_t swap_endian(uint32_t x) {
    return ((x & 0xFF) << 24) | ((x & 0xFF00) << 8) | 
           ((x & 0xFF0000) >> 8) | ((x & 0xFF000000) >> 24);
}

int sha256_init(sha256_ctx_t *ctx) {
    if (!ctx) return CRYPTO_ERROR_INVALID_PARAM;
    
    // 初始化状态
    for (int i = 0; i < 8; i++) {
        ctx->state[i] = H0[i];
    }
    
    // 初始化计数器
    ctx->count[0] = 0;
    ctx->count[1] = 0;
    
    return CRYPTO_SUCCESS;
}

int sha256_update(sha256_ctx_t *ctx, const uint8_t *data, size_t len) {
    if (!ctx || !data) return CRYPTO_ERROR_INVALID_PARAM;
    
    uint32_t i, index, part_len;
    
    // 计算当前缓冲区中的字节数
    index = (ctx->count[0] >> 3) & 0x3F;
    
    // 更新位计数
    if ((ctx->count[0] += (len << 3)) < (len << 3)) {
        ctx->count[1]++;
    }
    ctx->count[1] += (len >> 29);
    
    part_len = 64 - index;
    
    // 如果数据足够填满缓冲区
    if (len >= part_len) {
        memcpy(&ctx->buffer[index], data, part_len);
        sha256_transform(ctx, ctx->buffer);
        
        for (i = part_len; i + 63 < len; i += 64) {
            sha256_transform(ctx, &data[i]);
        }
        index = 0;
    } else {
        i = 0;
    }
    
    // 保存剩余数据
    memcpy(&ctx->buffer[index], &data[i], len - i);
    
    return CRYPTO_SUCCESS;
}

int sha256_final(sha256_ctx_t *ctx, uint8_t *digest) {
    if (!ctx || !digest) return CRYPTO_ERROR_INVALID_PARAM;
    
    uint8_t finalcount[8];
    uint32_t i;
    
    // 将位计数转换为字节序
    for (i = 0; i < 8; i++) {
        finalcount[i] = (uint8_t)((ctx->count[(i >= 4 ? 0 : 1)] >> ((3 - (i & 3)) * 8)) & 255);
    }
    
    // 添加填充
    sha256_update(ctx, (uint8_t*)"\200", 1);
    while ((ctx->count[0] & 1023) != 896) {
        sha256_update(ctx, (uint8_t*)"\0", 1);
    }
    sha256_update(ctx, finalcount, 8);
    
    // 生成最终哈希值
    for (i = 0; i < 32; i++) {
        digest[i] = (uint8_t)((ctx->state[i >> 2] >> ((3 - (i & 3)) * 8)) & 255);
    }
    
    return CRYPTO_SUCCESS;
}

int sha256_hash(const uint8_t *data, size_t len, uint8_t *digest) {
    if (!data || !digest) return CRYPTO_ERROR_INVALID_PARAM;
    
    sha256_ctx_t ctx;
    int ret = sha256_init(&ctx);
    if (ret != CRYPTO_SUCCESS) return ret;
    
    ret = sha256_update(&ctx, data, len);
    if (ret != CRYPTO_SUCCESS) return ret;
    
    return sha256_final(&ctx, digest);
}

// SHA256 变换函数
void sha256_transform(sha256_ctx_t *ctx, const uint8_t *data) {
    uint32_t a, b, c, d, e, f, g, h, i, j, t1, t2;
    uint32_t m[64];
    
    // 准备消息调度
    for (i = 0, j = 0; i < 16; ++i, j += 4) {
        m[i] = (data[j] << 24) | (data[j + 1] << 16) | (data[j + 2] << 8) | (data[j + 3]);
    }
    
    for (; i < 64; ++i) {
        m[i] = gamma1(m[i - 2]) + m[i - 7] + gamma0(m[i - 15]) + m[i - 16];
    }
    
    // 初始化工作变量
    a = ctx->state[0];
    b = ctx->state[1];
    c = ctx->state[2];
    d = ctx->state[3];
    e = ctx->state[4];
    f = ctx->state[5];
    g = ctx->state[6];
    h = ctx->state[7];
    
    // 主循环
    for (i = 0; i < 64; ++i) {
        t1 = h + sigma1(e) + ch(e, f, g) + K[i] + m[i];
        t2 = sigma0(a) + maj(a, b, c);
        h = g;
        g = f;
        f = e;
        e = d + t1;
        d = c;
        c = b;
        b = a;
        a = t1 + t2;
    }
    
    // 计算中间哈希值
    ctx->state[0] += a;
    ctx->state[1] += b;
    ctx->state[2] += c;
    ctx->state[3] += d;
    ctx->state[4] += e;
    ctx->state[5] += f;
    ctx->state[6] += g;
    ctx->state[7] += h;
} 